import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS

# --- 环境准备 ---
print("正在从 env.example 加载环境变量...")
load_dotenv(dotenv_path="env.example")
api_key = os.getenv("OPENAI_API_KEY")
# 注意：RAG 中我们需要两种模型
# 1. 用于生成回答的 LLM (比如 qwen-plus-latest)
# 2. 用于生成向量的 Embedding Model (我们也从配置文件读取)
# 因此，我们不再从环境变量读取模型名称，而是直接在代码中指定。
llm_model_name = os.getenv("OPENAI_MODEL", "qwen-plus-latest")
embedding_model_name = os.getenv("OPENAI_EMBEDDING_MODEL", "text-embedding-v1")

if not api_key:
    print("错误：请先在 env.example 文件中设置你的 OPENAI_API_KEY。")
else:
    print("环境变量加载成功！")
    # --- RAG 核心流程 ---

    # 1. 加载 (Load)
    # 使用 TextLoader 来加载我们的知识库文件
    loader = TextLoader("my_knowledge.txt", encoding="utf-8")
    docs = loader.load()
    print("1. 文档加载完成。")
    # print("加载的文档内容:", docs[0].page_content[:100]) # 打印前100个字符看看

    # 2. 切分 (Split)
    # 我们使用 RecursiveCharacterTextSplitter，这是一个通用的文本切分器
    # chunk_size 定义了每个文本块的最大长度（按字符数计）
    # chunk_overlap 定义了相邻文本块之间的重叠字符数，这有助于保持上下文的连续性
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
    splits = text_splitter.split_documents(docs)
    print("2. 文档切分完成。")
    # print("切分后的第一个块:", splits[0])

    # 3. 存储 (Store)
    # 我们需要一个"嵌入模型"来将文本转换为向量
    # 因为我们使用的是阿里云DashScope，所以我们应该使用专门的DashScopeEmbeddings类。
    # 它知道如何正确地与DashScope API通信。我们把从环境变量中读到的key传给它。
    embeddings = DashScopeEmbeddings(
        model=embedding_model_name,
        dashscope_api_key=api_key
    )
    
    # 我们使用 FAISS 这个库在本地内存中创建一个向量数据库
    # from_documents 方法会自动为所有切分后的文本块生成向量并存储起来
    vectorstore = FAISS.from_documents(documents=splits, embedding=embeddings)
    print("3. 向量数据库创建完成。")

    # 4. 检索 (Retrieve)
    # 我们从向量数据库创建一个"检索器"对象
    # 它的任务就是根据一个查询，返回最相关的文档
    retriever = vectorstore.as_retriever()
    print("4. 检索器已准备就绪。")

    # 5. 生成 (Generate)
    # a. 创建提示模板
    # 我们需要给模型一个明确的指令：根据我提供的上下文来回答问题
    # `{context}`: 这个占位符将用来存放我们从数据库中检索到的相关信息
    # `{question}`: 这个占位符将用来存放用户的原始问题
    prompt = ChatPromptTemplate.from_template(
        """
        你是一个专业的问答助手。请仅仅依据我提供的以下上下文信息来回答问题。
        如果你在上下文中找不到答案，就直接说"我不知道"。

        上下文:
        {context}

        问题:
        {question}
        """
    )

    # b. 创建 RAG 链
    # 这条链的逻辑完美地体现了 RAG 的流程
    rag_chain = (
        # RunnablePassthrough.assign(...) 的作用是创建一个包含 context 和 question 的字典
        # context 的内容来自于 retriever 的输出
        # question 的内容来自于原始的输入
        RunnablePassthrough.assign(context=(lambda x: x["question"]) | retriever)
        | prompt
        | ChatOpenAI(model=llm_model_name)
        | StrOutputParser()
    )
    print("5. RAG 链创建完成！")

    # 6. 开始提问！
    print("\n--- 开始向你的知识库提问 ---")
    question_1 = "DocuMind 是谁开发的？"
    response_1 = rag_chain.invoke({"question": question_1})
    print(f"[问题]: {question_1}")
    print(f"[回答]: {response_1}")

    question_2 = "DocuMind 有什么主要功能？"
    response_2 = rag_chain.invoke({"question": question_2})
    print(f"\n[问题]: {question_2}")
    print(f"[回答]: {response_2}")
    
    # 这是一个"陷阱"问题，答案不在我们的知识库里
    question_3 = "DocuMind 支持处理视频文件吗？"
    response_3 = rag_chain.invoke({"question": question_3})
    print(f"\n[问题]: {question_3}")
    print(f"[回答]: {response_3}") # AI 应该回答 "我不知道"

    print("\n--- 提问结束 ---") 